# coding:utf-8
import pandas as pd
from xtquant import xtdata, xtconstant, xttrader
from Config import *
import datetime
import numpy as np
import numpy
import math

# 定义一个全局变量，设定要操作的股票
# C.stock_list = C.get_stock_list_in_sector("沪深300")  # 获取沪深300股票列表
pct = 0.5
code_list = ['159915.SZ', '510300.SH', '510500.SH', '513100.SH', '518880.SH', '513520.SH', '159509.SZ']
# code_list = ['518880.SH']
your_fund = 7000
# 北向净流入数据
Northbound = []
# 针对300etf优化选值,初始为15
Northbound_val = 15
north_money = 0

# 动量轮动参数
stock_num = 1  # 买入评分最高的前stock_num只股票
momentum_day = 25  # 最新动量参考最近momentum_day的
# rsrs择时参数
ref_stock = '000300.SH'  # 用ref_stock做择时计算的基础数据
N = 18  # 计算最新斜率slope，拟合度r2参考最近N(18)天
M = 600  # 计算最新标准分zscore，rsrs_score参考最近M(600)天
score_threshold = 0.7  # rsrs标准分指标阈值
# ma择时参数
mean_day = 20  # 计算结束ma收盘价，参考最近mean_day
mean_diff_day = 5  # 计算初始ma收盘价，参考(mean_day + mean_diff_day)天前，窗口为mean_diff_day的一段时间


# 除去回测第一天的slope，避免运行时重复加入
def callback_func(data):
    pass


def clock_trade():
    current_time = datetime.datetime.now().strftime('%H:%M:%S')
    if current_time > datetime.time(8, 30, 50).strftime('%H:%M:%S') and current_time < datetime.time(8, 30, 55).strftime('%H:%M:%S'):
        my_trade()


def mius_and_get_start_date(n_days):
    end_date = datetime.datetime.now()
    start_date = end_date - datetime.timedelta(days=n_days)
    return start_date.strftime('%Y%m%d')


# 获取n天前的股票收盘信息
def get_history_data(n_days, code, fields, period_str):
    xtdata.download_history_data(code, period='1d', start_time=mius_and_get_start_date(n_days))
    data = xtdata.get_market_data(field_list=fields, stock_list=[code], period=period_str,
                                  start_time=mius_and_get_start_date(n_days))
    # print(data)
    return data


# 1-1 选股模块-动量因子轮动
# 基于股票年化收益和判定系数打分,并按照分数从大到小排名

def get_rank(code_list):
    score_list = []
    for stock in code_list:
        data = get_history_data(n_days=momentum_day, code=stock, fields=['close'], period_str='1d')
        # 收盘价List<List>
        y = data['log'] = np.log(data['close'].values[0])
        x = data['num'] = np.arange(data['log'].size)
        slope, intercept = np.polyfit(x, y, 1)
        # print(slope, intercept)
        annualized_returns = math.pow(math.exp(slope), 250) - 1
        r_squared = 1 - (sum((y - (slope * x + intercept)) ** 2) / ((len(y) - 1) * np.var(y, ddof=1)))
        score = annualized_returns * r_squared
        score_list.append(score)
    stock_dict = dict(zip(code_list, score_list))
    sort_list = sorted(stock_dict.items(), key=lambda item: item[1], reverse=True)  # True为降序
    res_list = []
    for i in range((len(code_list))):
        res_list.append(sort_list[i][0])
    rank_stock = res_list[0:stock_num]
    # print(res_list[0:5])
    return rank_stock


# 2-1 择时模块-计算线性回归统计值
# 对输入的自变量每日最低价x(series)和因变量每日最高价y(series)建立OLS回归模型,返回元组(截距,斜率,拟合度)
def get_ols(x, y):
    if len(x) != 0 and len(y) != 0:
        slope, intercept = np.polyfit(x, y, 1)
        r2 = 1 - (sum((y - (slope * x + intercept)) ** 2) / ((len(y) - 1) * np.var(y, ddof=1)))
        return (intercept, slope, r2)
    return []


# 2-2 择时模块-设定初始斜率序列
# 通过前M日最高最低价的线性回归计算初始的斜率,返回斜率的列表
def initial_slope_series():
    nm_days = N + M
    data = get_history_data(n_days=nm_days, code=ref_stock, period_str='1d',
                            fields=['high', 'low'])
    results = []
    for i in range(M):
        low_data = data['low'].values[0][i:i + N]
        high_data = data['high'].values[0][i:i + N]
        ols_result = get_ols(low_data, high_data)
        if len(ols_result) >= 1:
            results.append(ols_result[1])
    return results


# 2-3 择时模块-计算标准分
# 通过斜率列表计算并返回截至回测结束日的最新标准分
def get_zscore(slope_series):
    mean = np.mean(slope_series)
    std = np.std(slope_series)
    return (slope_series[-1] - mean) / std


# 2-4 择时模块-计算综合信号
# 1.获得rsrs与MA信号,rsrs信号算法参考优化说明，MA信号为一段时间两个端点的MA数值比较大小
# 2.信号同时为True时返回买入信号，同为False时返回卖出信号，其余情况返回持仓不变信号
def get_timing_signal():
    # 修改
    slope_series = initial_slope_series()[:-1]
    # print(slope_series)
    # 计算MA信号
    close_data = get_history_data(n_days=mean_day + mean_diff_day, code=ref_stock, period_str='1d', fields=['close'])
    today_MA = close_data['close'].values[0][mean_diff_day:].mean()
    before_MA = close_data['close'].values[0][:-mean_diff_day].mean()
    # print("today_MA:{}".format(today_MA))
    # print("before_MA:{}".format(before_MA))
    # 计算rsrs信号
    high_low_data = get_history_data(n_days=N, code=ref_stock, period_str='1d', fields=['high', 'low'])
    intercept, slope, r2 = get_ols(high_low_data['low'].values[0], high_low_data['high'].values[0])
    # print("intercept:{}".format(intercept))
    slope_series.append(slope)
    rsrs_score = get_zscore(slope_series[-M:]) * r2
    # 综合判断所有信号
    print(rsrs_score, score_threshold)
    if rsrs_score > score_threshold and today_MA > before_MA:
        # print('BUY')
        return "BUY"
    elif rsrs_score < -score_threshold and today_MA < before_MA:
        # print('SELL')
        return "SELL"
    else:
        # print('KEEP')
        return "KEEP"


# 2-5 择时模块-对300ETF进行北向净值判断
# 获取前一交易日北向资金流向，参数exp:2021-10-21
def get_finance(context):
    yesterday = context.previous_date
    # print(today.date(),'@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@')

    # 20161205前北向资金为沪股通，20161205深股通数据
    n_sh = finance.run_query(query(finance.STK_ML_QUOTA).filter(finance.STK_ML_QUOTA.day <= yesterday,
                                                                finance.STK_ML_QUOTA.link_id == 310001).order_by(
        finance.STK_ML_QUOTA.day.desc()).limit(10))
    if n_sh.empty:
        return 0
    n_sz = finance.run_query(query(finance.STK_ML_QUOTA).filter(finance.STK_ML_QUOTA.day <= yesterday,
                                                                finance.STK_ML_QUOTA.link_id == 310002).order_by(
        finance.STK_ML_QUOTA.day.desc()).limit(10))
    total_net_in = 0
    if n_sz.empty:
        for i in range(0, 1):
            sh_in = n_sh['buy_amount'][i] - n_sh['sell_amount'][i]
            total_net_in += sh_in
    else:
        for i in range(0, 1):
            sh_in = n_sh['buy_amount'][i] - n_sh['sell_amount'][i]
            sz_in = n_sz['buy_amount'][i] - n_sz['sell_amount'][i]
            amount = sh_in + sz_in
            total_net_in += amount

    north_money = total_net_in
    Northbound.append(g.north_money)
    new = sorted(Northbound)
    # 中位数
    if (numpy.size(Northbound) > 100):  # 前100天默认为15e且买入val>0
        Northbound_val = new[numpy.size(new) // 2]
        if Northbound_val < 5:
            Northbound_val = 5


def get_close_price(code):
    data = get_history_data(n_days=2, code=code, fields=['close'], period_str='1d')
    return data['close'].values[0]


# 获取持仓列表和数量,key:code,val:volume
def get_hold_stock_list():
    positions = xt_trader.query_stock_positions(account_putong)
    hold_dic = {}
    if len(positions) == 0:
        return hold_dic
    for pos in positions:
        hold_dic[pos.stock_code] = pos.volume
    return hold_dic


# 交易模块-择时交易
# 结合择时模块综合信号进行交易
def my_trade():
    check_out_list = get_rank(code_list)
    print('今日自选股:{}'.format(check_out_list))

    timing_signal = get_timing_signal()
    print(timing_signal)

    last_price = get_close_price(check_out_list[0])
    # 下单前准备工作
    msg = {}
    msg['pct'] = pct
    msg['code'] = check_out_list[0]
    msg['last_price'] = last_price.item()
    msg['action'] = timing_signal

    # 下单函数
    order_handle(msg)


# 下单函数
def order_handle(msg):
    # 根据msg给交易端下单
    stock_code = msg['code']
    # 下单比例
    order_rate = msg['pct']
    last_price = msg['last_price']
    # print(last_price)

    if msg['action'] == 'BUY':
        # print(msg['action'])
        # 计算买入价格，默认1.015价格。
        adjusted_price = round(last_price * 1.01, 2)
        # 计算股数
        # 下单资金
        use_money = your_fund * order_rate
        stock_quantity = use_money / adjusted_price  # 符合条件的5000配置
        stock_quantity -= stock_quantity % 100  # 取整百
        stock_quantity = int(stock_quantity)  # 将结果转换为整数
        # print(stock_quantity)
        # stock_quantity = 100
        fix_result_order_id = xt_trader.order_stock(account_putong, stock_code, xtconstant.STOCK_BUY,
                                                    stock_quantity,
                                                    xtconstant.FIX_PRICE, adjusted_price, '聚宽策略', '买入')
        print("订单号:", fix_result_order_id)
        print("CODE", stock_code)
        print("比例", order_rate)
        print("价格:", last_price * 1.015)

    if msg['action'] == 'SELL':
        print(msg['action'])
        # 计算买入价格，默认1.015价格。
        adjusted_price = round(last_price * 0.985, 2)
        # 计算股数
        # 下单资金
        code_volume_dict = get_hold_stock_list()
        # 全仓卖出
        stock_quantity = code_volume_dict.get(stock_code)
        # print(stock_quantity)
        # stock_quantity = 100
        fix_result_order_id = xt_trader.order_stock(account_putong, stock_code, xtconstant.STOCK_SELL, stock_quantity,
                                                    xtconstant.FIX_PRICE, adjusted_price, '聚宽策略', '卖出')
        print("订单号:", fix_result_order_id)
        print("CODE", stock_code)
        print("比例", order_rate)
        print("价格:", last_price * 0.985)
    # TODO keep如何动态调整仓位
    if msg['action'] == 'KEEP':
        print(msg['action'])
        # 获取持仓列表
        code_volume_dict = get_hold_stock_list()
        # 已经持有，return，否则继续买入持仓
        if stock_code in code_volume_dict.keys():
            return
        # 计算买入价格，默认1.015价格。
        adjusted_price = round(last_price * 1.015, 2)
        # 计算股数
        # 下单资金
        use_money = your_fund * order_rate
        stock_quantity = use_money / adjusted_price  # 符合条件的5000配置
        stock_quantity -= stock_quantity % 100  # 取整百
        stock_quantity = int(stock_quantity)  # 将结果转换为整数
        print(stock_quantity)
        # stock_quantity = 100
        fix_result_order_id = xt_trader.order_stock(account_putong, stock_code, xtconstant.STOCK_BUY,
                                                    stock_quantity,
                                                    xtconstant.FIX_PRICE, adjusted_price, '聚宽策略', '买入')
        print("订单号:", fix_result_order_id)
        print("CODE", stock_code)
        print("比例", order_rate)
        print("价格:", last_price * 1.015)


# ==========订阅信息==========

xtdata.subscribe_whole_quote(code_list, callback=callback_func)

# 每天监听8：50的下单信号
clock_trade()

xtdata.run()
